package adkr

import (
	"ebbflow/pkg/crypto"
	"encoding/json"
	"log"
	"math/big"
	"strconv"
	"sync"
)

// STUB for ProofOfKnowledge.
func ProofOfKnowledge(base crypto.G1, pub crypto.G1, sec crypto.ZR) (PoK, error) {
	var c, z crypto.ZR
	c.SetRandom()
	z.SetRandom()
	return PoK{C: c, Z: z}, nil
}

// STUB for PokVerify.
func PokVerify(pub crypto.G1, base crypto.G1, proof PoK) bool {
	return true
}

// acssResult is a helper struct to pair an ACSS output with its dealer ID.
type acssResult struct {
	dealerID int
	output   ACSSOutput
}

// ADKR is the main function that orchestrates the entire protocol.
func ADKR(
	sid string, pid, n, f int, g, h *crypto.G1, ePKs []crypto.G1, my_eSK crypto.ZR,
	msgInCh <-chan Message, msgOutCh chan<- Message,
) {
	log.Printf("ADKR-SID:%s PID:%d Starting main protocol logic...", sid, pid)

	// --- 1. Channel Setup and Dispatcher ---
	acssInChs := make([]chan Message, n)
	randexInCh := make(chan Message, 100)
	keyInCh := make(chan Message, 100)
	mvbaInCh := make(chan Message, 100)
	for i := 0; i < n; i++ {
		acssInChs[i] = make(chan Message, 100)
	}

	go func() {
		for msg := range msgInCh {
			switch msg.Tag {
			case ADKR_ACSS:
				if msg.Round >= 0 && msg.Round < n {
					acssInChs[msg.Round] <- msg
				}
			case ADKR_MVBA:
				mvbaInCh <- msg
			case ADKG_RANDEX:
				randexInCh <- msg
			case ADKG_KEY:
				keyInCh <- msg
			}
		}
	}()

	// --- 2. ACSS Phase ---
	log.Printf("ADKR-SID:%s PID:%d Starting %d ACSS instances...", sid, pid, n)
	acssOutputChs := make([]chan ACSSOutput, n)
	allAcssOutputs := make(map[int]ACSSOutput)
	acssInputChs := make([]chan [2]crypto.ZR, n)

	for i := 0; i < n; i++ {
		acssInputChs[i] = make(chan [2]crypto.ZR, 1)
		acssOutputChs[i] = make(chan ACSSOutput, 1)
		var inputChForACSS chan [2]crypto.ZR
		if i == pid {
			inputChForACSS = acssInputChs[i]
		}
		go CompleteSecretSharing(
			sid+"ACSS"+strconv.Itoa(i), pid, n, f, nil, n, f, nil, g, h, i,
			ePKs, my_eSK, inputChForACSS, acssOutputChs[i], acssInChs[i], msgOutCh,
		)
	}

	go func() {
		secretA := crypto.ZRFromInt(int64(123 + pid))
		secretB := crypto.ZRFromInt(int64(456 + pid))
		acssInputChs[pid] <- [2]crypto.ZR{secretA, secretB}
		log.Printf("ADKR-SID:%s PID:%d Provided secret to own ACSS instance.", sid, pid)
	}()

	acssResultsCh := make(chan acssResult, n)
	var wg sync.WaitGroup
	wg.Add(n)
	for i := 0; i < n; i++ {
		go func(dealerID int) {
			defer wg.Done()
			output := <-acssOutputChs[dealerID]
			allAcssOutputs[dealerID] = output
			acssResultsCh <- acssResult{dealerID: dealerID, output: output}
		}(i)
	}
	go func() {
		wg.Wait()
		close(acssResultsCh)
	}()
	completedDealers := make(map[int]struct{})
	for result := range acssResultsCh {
		log.Printf("ADKR-SID:%s PID:%d Collected output from ACSS for dealer %d", sid, pid, result.dealerID)
		completedDealers[result.dealerID] = struct{}{}
	}

	// --- 3. MVBA Phase ---
	requiredOutputs := 2*f + 1
	proposalSet := make([]int, 0, requiredOutputs)
	for id := range completedDealers {
		if len(proposalSet) < requiredOutputs {
			proposalSet = append(proposalSet, id)
		}
	}
	log.Printf("ADKR-SID:%s PID:%d Proposing %v for MVBA.", sid, pid, proposalSet)

	mvbaInputCh := make(chan []byte, 1)
	mvbaOutputCh := make(chan []byte, 1)
	localSetCh := make(chan int, n)

	predicate := func(value []byte, localSet map[int]struct{}) bool {
		var receivedSet []int
		if err := json.Unmarshal(value, &receivedSet); err != nil {
			return false
		}
		if len(receivedSet) < requiredOutputs {
			return false
		}
		for _, id := range receivedSet {
			if _, ok := localSet[id]; !ok {
				return false
			}
		}
		return true
	}
	go func() {
		for id := range completedDealers {
			localSetCh <- id
		}
		close(localSetCh)
	}()

    // 优先使用外部注册的阈值 BLS 配置；若未注册，则回退为本地演示生成
    pks, thPK := GetThresholdBLSPKs()
    var mySK crypto.ZR
    if GetMyBLSSK() != nil { mySK = *GetMyBLSSK() }
    if len(pks) != n || thPK == nil || GetMyBLSSK() == nil {
        genPK, pksGen, skShares := crypto.GenerateThresholdKeys(n, f+1)
        pks = pksGen
        mySK = skShares[pid]
        thPK = genPK
    }
    go TThresholdMVBA(sid+"MVBA", pid, n, f, nil, g, mvbaInputCh, mvbaOutputCh, mvbaInCh, msgOutCh, nil, localSetCh, predicate, pks, mySK, thPK)
	proposalBytes, _ := json.Marshal(proposalSet)
	mvbaInputCh <- proposalBytes

	log.Printf("ADKR-SID:%s PID:%d Waiting for MVBA to decide...", sid, pid)
	finalDecisionBytes := <-mvbaOutputCh
	var finalSet []int
	json.Unmarshal(finalDecisionBytes, &finalSet)
	log.Printf("ADKR-SID:%s PID:%d MVBA DECIDED on set: %v!", sid, pid, finalSet)

	// --- 4. Key Refresh Computation ---
	log.Printf("ADKR-SID:%s PID:%d Proceeding to Key Refresh: Random Extraction.", sid, pid)
	finalSetMap := make(map[int]struct{})
	for _, id := range finalSet {
		finalSetMap[id] = struct{}{}
	}
	secretsA := make([]crypto.ZR, n)
	randomnessA := make([]crypto.ZR, n)
	secretsB := make([]crypto.ZR, n)
	randomnessB := make([]crypto.ZR, n)
	for i := 0; i < n; i++ {
		if out, ok := allAcssOutputs[i]; ok && (len(finalSet) == 0 || contains(finalSet, i)) {
			secretsA[i] = out.ShareA
			randomnessA[i] = out.RandomnessA
			secretsB[i] = out.ShareB
			randomnessB[i] = out.RandomnessB
		} else {
			secretsA[i].SetZero()
			randomnessA[i].SetZero()
			secretsB[i].SetZero()
			randomnessB[i].SetZero()
		}
	}
	mat1 := make([][]crypto.ZR, f+1)
	for i := 0; i <= f; i++ {
		mat1[i] = make([]crypto.ZR, n)
		for j := 0; j < n; j++ {
			var exp big.Int
			exp.SetInt64(int64(i))
			base := crypto.ZRFromInt(int64(j + 1))
			mat1[i][j].Exp(base, &exp)
		}
	}
	mat2 := make([][]crypto.ZR, f)
	for i := 0; i < f; i++ {
		mat2[i] = make([]crypto.ZR, n)
		for j := 0; j < n; j++ {
			var exp big.Int
			exp.SetInt64(int64(i))
			base := crypto.ZRFromInt(int64(j + 1))
			mat2[i][j].Exp(base, &exp)
		}
	}
	zShares := make([]crypto.ZR, 2*f+1)
	rShares := make([]crypto.ZR, 2*f+1)
	for i := 0; i <= f; i++ {
		zShares[i], _ = crypto.InnerProduct(mat1[i], secretsA)
		rShares[i], _ = crypto.InnerProduct(mat1[i], randomnessA)
	}
	for i := 0; i < f; i++ {
		zShares[f+1+i], _ = crypto.InnerProduct(mat2[i], secretsB)
		rShares[f+1+i], _ = crypto.InnerProduct(mat2[i], randomnessB)
	}
	log.Printf("ADKR-SID:%s PID:%d Computed %d new shares.", sid, pid, len(zShares))

	// --- 5. Key Refresh Finalization ---
	log.Printf("ADKR-SID:%s PID:%d Broadcasting new shares (RANDEX)...", sid, pid)
	var payloadBytes []byte
	for i := 0; i < n; i++ {
		s_i := crypto.EvalPolynomial(zShares, int64(i+1))
		r_i := crypto.EvalPolynomial(rShares, int64(i+1))
		randexPayload := RandexPayload{S: s_i, R: r_i}
		payloadBytes, _ = json.Marshal(randexPayload)
		msgOutCh <- Message{Tag: ADKG_RANDEX, To: i, Payload: payloadBytes}
	}

	randexShares := make([][2]interface{}, 0, f+1)
	for i := 0; i < f+1; i++ {
		msg := <-randexInCh
		var payload RandexPayload
		json.Unmarshal(msg.Payload, &payload)
		randexShares = append(randexShares, [2]interface{}{int64(msg.Sender + 1), payload.S})
	}
	log.Printf("ADKR-SID:%s PID:%d Received %d RANDEX shares.", sid, pid, len(randexShares))

	zi, _ := crypto.InterpolateAtX(randexShares, 0)
	newPKShare := crypto.MulG1(g, &zi)
	pok, _ := ProofOfKnowledge(newPKShare, *g, zi)
	keyPayload := KeyPayload{Sgz: newPKShare, Pi: pok}
	payloadBytes, _ = json.Marshal(keyPayload)
	msgOutCh <- Message{Tag: ADKG_KEY, To: -1, Payload: payloadBytes}
	log.Printf("ADKR-SID:%s PID:%d Reconstructed own key and broadcasted KEY message.", sid, pid)

	verifiedPKs := 0
	for msg := range keyInCh {
		var payload KeyPayload
		json.Unmarshal(msg.Payload, &payload)
		if PokVerify(payload.Sgz, *g, payload.Pi) {
			verifiedPKs++
			log.Printf("ADKR-SID:%s PID:%d Verified KEY message from Node %d.", sid, pid, msg.Sender)
		} else {
			log.Printf("ADKR-SID:%s PID:%d INVALID KEY message from Node %d.", sid, pid, msg.Sender)
		}
		if verifiedPKs >= n-f {
			break
		}
	}

	log.Printf("ADKR-SID:%s PID:%d Protocol finished successfully! All stages complete.", sid, pid)
	<-make(chan struct{})
}

func contains(s []int, e int) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}
